# -*- coding: utf-8 -*- 
# @Time : 2022/4/2 10:07 
# @Author : zzuxyj 
# @File : 06-nn_conv_simulate02.py

"""
使用nn.Moudle 构造一个神经网络，并使用卷积进行计算

会计算卷积之后的大小
"""
import torch
import torchvision
from torch import nn
from torch.nn import Conv2d
from torch.utils.data import DataLoader

# 数据集和数据加载
from torch.utils.tensorboard import SummaryWriter

trainSet = torchvision.datasets.CIFAR100(root="../dataset/CIFAR100/", download=True, transform=torchvision.transforms.ToTensor() ,train= True)
testSet = torchvision.datasets.CIFAR100(root="../dataset/CIFAR100/", download=True, transform=torchvision.transforms.ToTensor() ,train= False)
dataloader = DataLoader(dataset=testSet , batch_size=64 , drop_last=False , shuffle=True )

# tensorboard
writer = SummaryWriter("logs06")

# 一个卷积的神经网络模型
class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = Conv2d(in_channels=3, out_channels=6, kernel_size=3, stride=1, padding=0)
        # (32 - 3 + 1 ) / 1 = 30
    def forward(self, tensor):
        x = self.conv1(tensor)
        return x;

# 创建对象
model = Model()

# 遍历数据集
step = 0
for data in dataloader:
    imgs , target = data
    writer.add_images("CIFAR100" , imgs , global_step=step)
    output = model(imgs)
    output = torch.reshape(output , (-1,3,30,30))
    print(output.shape)
    writer.add_images("train_CIFAR100" , output , global_step=step)
    step+=1

#关闭流
writer.close()